import tkinter as tk
from tkinter import ttk, messagebox, filedialog
import os
import shutil
import json
from pathlib import Path
import threading
from datetime import datetime
import sys
from config import get_file_categories, get_setting, set_setting

class DesktopHelper:
    def __init__(self):
        self.root = tk.Tk()
        self.root.title("桌面助手 - Desktop Helper")
        self.root.geometry("800x600")
        self.root.configure(bg='#f0f0f0')
        
        # 设置窗口图标和属性
        self.root.resizable(True, True)
        self.root.minsize(600, 400)
        
        # 从配置文件加载文件分类规则
        self.file_categories = get_file_categories()
        
        # 获取桌面路径
        self.desktop_path = self.get_desktop_path()
        self.storage_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "桌面收纳空间")
        self.backup_path = os.path.join(self.storage_path, "备份记录")
        
        # 确保存储目录存在
        os.makedirs(self.storage_path, exist_ok=True)
        os.makedirs(self.backup_path, exist_ok=True)
        
        self.setup_ui()
        
    def get_desktop_path(self):
        """获取桌面路径"""
        return os.path.join(os.path.expanduser("~"), "Desktop")
    
    def setup_ui(self):
        """设置用户界面"""
        # 主标题
        title_frame = tk.Frame(self.root, bg='#2c3e50', height=80)
        title_frame.pack(fill='x', pady=0)
        title_frame.pack_propagate(False)
        
        title_label = tk.Label(title_frame, text="🖥️ 桌面助手", font=('微软雅黑', 20, 'bold'), 
                              fg='white', bg='#2c3e50')
        title_label.pack(expand=True)
        
        subtitle_label = tk.Label(title_frame, text="让您的桌面更加整洁有序", font=('微软雅黑', 12), 
                                 fg='#ecf0f1', bg='#2c3e50')
        subtitle_label.pack()
        
        # 创建notebook（选项卡）
        self.notebook = ttk.Notebook(self.root)
        self.notebook.pack(fill='both', expand=True, padx=10, pady=10)
        
        # 桌面整理选项卡
        self.setup_organize_tab()
        
        # 设置选项卡
        self.setup_settings_tab()
        
        # 关于选项卡
        self.setup_about_tab()
        
    def setup_organize_tab(self):
        """设置桌面整理选项卡"""
        organize_frame = ttk.Frame(self.notebook)
        self.notebook.add(organize_frame, text="桌面整理")
        
        # 创建滚动区域
        canvas = tk.Canvas(organize_frame, bg='white')
        scrollbar = ttk.Scrollbar(organize_frame, orient="vertical", command=canvas.yview)
        scrollable_frame = ttk.Frame(canvas)
        
        scrollable_frame.bind(
            "<Configure>",
            lambda e: canvas.configure(scrollregion=canvas.bbox("all"))
        )
        
        canvas.create_window((0, 0), window=scrollable_frame, anchor="nw")
        canvas.configure(yscrollcommand=scrollbar.set)
        
        # 桌面状态区域
        status_frame = tk.LabelFrame(scrollable_frame, text="桌面状态", font=('微软雅黑', 12, 'bold'), 
                                   bg='white', fg='#2c3e50')
        status_frame.pack(fill='x', padx=20, pady=10)
        
        self.status_text = tk.Text(status_frame, height=8, font=('微软雅黑', 10), 
                                  state='disabled', bg='#f8f9fa')
        self.status_text.pack(fill='x', padx=10, pady=10)
        
        # 控制按钮区域
        control_frame = tk.LabelFrame(scrollable_frame, text="整理控制", font=('微软雅黑', 12, 'bold'), 
                                    bg='white', fg='#2c3e50')
        control_frame.pack(fill='x', padx=20, pady=10)
        
        button_frame = tk.Frame(control_frame, bg='white')
        button_frame.pack(pady=15)
        
        # 扫描桌面按钮
        scan_btn = tk.Button(button_frame, text="🔍 扫描桌面", font=('微软雅黑', 12, 'bold'), 
                           bg='#3498db', fg='white', relief='flat', padx=20, pady=10,
                           command=self.scan_desktop)
        scan_btn.pack(side='left', padx=10)
        
        # 一键整理按钮
        organize_btn = tk.Button(button_frame, text="✨ 一键整理", font=('微软雅黑', 12, 'bold'), 
                               bg='#e74c3c', fg='white', relief='flat', padx=20, pady=10,
                               command=self.organize_desktop)
        organize_btn.pack(side='left', padx=10)
        
        # 还原桌面按钮
        restore_btn = tk.Button(button_frame, text="↩️ 还原桌面", font=('微软雅黑', 12, 'bold'), 
                              bg='#f39c12', fg='white', relief='flat', padx=20, pady=10,
                              command=self.restore_desktop)
        restore_btn.pack(side='left', padx=10)
        
        # 打开收纳空间按钮
        open_storage_btn = tk.Button(button_frame, text="📁 打开收纳空间", font=('微软雅黑', 12, 'bold'), 
                                   bg='#27ae60', fg='white', relief='flat', padx=20, pady=10,
                                   command=self.open_storage)
        open_storage_btn.pack(side='left', padx=10)
        
        # 进度条
        self.progress_var = tk.DoubleVar()
        self.progress_bar = ttk.Progressbar(control_frame, variable=self.progress_var, 
                                          maximum=100, length=300)
        self.progress_bar.pack(pady=10)
        
        # 收纳空间概览
        storage_frame = tk.LabelFrame(scrollable_frame, text="收纳空间概览", font=('微软雅黑', 12, 'bold'), 
                                    bg='white', fg='#2c3e50')
        storage_frame.pack(fill='x', padx=20, pady=10)
        
        self.storage_tree = ttk.Treeview(storage_frame, columns=('count', 'size'), show='tree headings')
        self.storage_tree.heading('#0', text='分类')
        self.storage_tree.heading('count', text='文件数量')
        self.storage_tree.heading('size', text='总大小')
        self.storage_tree.pack(fill='x', padx=10, pady=10)
        
        canvas.pack(side="left", fill="both", expand=True)
        scrollbar.pack(side="right", fill="y")
        
        # 初始扫描
        self.scan_desktop()
        
    def setup_settings_tab(self):
        """设置选项卡"""
        settings_frame = ttk.Frame(self.notebook)
        self.notebook.add(settings_frame, text="设置")
        
        # 创建滚动区域
        canvas = tk.Canvas(settings_frame, bg='white')
        scrollbar = ttk.Scrollbar(settings_frame, orient="vertical", command=canvas.yview)
        scrollable_frame = ttk.Frame(canvas)
        
        scrollable_frame.bind(
            "<Configure>",
            lambda e: canvas.configure(scrollregion=canvas.bbox("all"))
        )
        
        canvas.create_window((0, 0), window=scrollable_frame, anchor="nw")
        canvas.configure(yscrollcommand=scrollbar.set)
        
        # 文件分类设置
        file_types_frame = tk.LabelFrame(scrollable_frame, text="文件分类设置", 
                                        font=('微软雅黑', 12, 'bold'))
        file_types_frame.pack(fill='x', padx=20, pady=10)
        
        # 显示当前分类规则
        for category, extensions in self.file_categories.items():
            category_frame = tk.Frame(file_types_frame, bg='white')
            category_frame.pack(fill='x', padx=10, pady=5)
            
            tk.Label(category_frame, text=f"{category}:", font=('微软雅黑', 10, 'bold'),
                    bg='white').pack(side='left')
            
            # 显示扩展名（只显示前几个）
            ext_text = ", ".join(extensions[:5])
            if len(extensions) > 5:
                ext_text += f" ... (共{len(extensions)}个)"
            
            tk.Label(category_frame, text=ext_text, font=('微软雅黑', 9),
                    bg='white', fg='#666').pack(side='left', padx=(10, 0))
        
        # 整理设置
        organize_frame = tk.LabelFrame(scrollable_frame, text="整理设置", 
                                     font=('微软雅黑', 12, 'bold'))
        organize_frame.pack(fill='x', padx=20, pady=10)
        
        # 自动备份设置
        self.auto_backup_var = tk.BooleanVar(value=get_setting("auto_backup", True))
        auto_backup_cb = tk.Checkbutton(organize_frame, text="整理前自动创建备份", 
                                       variable=self.auto_backup_var,
                                       font=('微软雅黑', 10),
                                       command=self.on_setting_changed)
        auto_backup_cb.pack(anchor='w', padx=10, pady=5)
        
        # 确认对话框设置
        self.confirm_organize_var = tk.BooleanVar(value=get_setting("confirm_before_organize", True))
        confirm_cb = tk.Checkbutton(organize_frame, text="整理前显示确认对话框", 
                                   variable=self.confirm_organize_var,
                                   font=('微软雅黑', 10),
                                   command=self.on_setting_changed)
        confirm_cb.pack(anchor='w', padx=10, pady=5)
        
        # 跳过系统文件设置
        self.skip_system_var = tk.BooleanVar(value=get_setting("skip_system_files", True))
        skip_system_cb = tk.Checkbutton(organize_frame, text="跳过系统文件和重要快捷方式", 
                                       variable=self.skip_system_var,
                                       font=('微软雅黑', 10),
                                       command=self.on_setting_changed)
        skip_system_cb.pack(anchor='w', padx=10, pady=5)
        
        # 高级设置
        advanced_frame = tk.LabelFrame(scrollable_frame, text="高级设置", 
                                     font=('微软雅黑', 12, 'bold'))
        advanced_frame.pack(fill='x', padx=20, pady=10)
        
        # 最大备份数量
        backup_count_frame = tk.Frame(advanced_frame, bg='white')
        backup_count_frame.pack(fill='x', padx=10, pady=5)
        
        tk.Label(backup_count_frame, text="最大备份数量:", font=('微软雅黑', 10),
                bg='white').pack(side='left')
        
        self.backup_count_var = tk.StringVar(value=str(get_setting("max_backup_count", 10)))
        backup_spinbox = tk.Spinbox(backup_count_frame, from_=1, to=50, width=10,
                                   textvariable=self.backup_count_var,
                                   font=('微软雅黑', 9),
                                   command=self.on_setting_changed)
        backup_spinbox.pack(side='left', padx=(10, 0))
        
        # 按钮区域
        button_frame = tk.Frame(scrollable_frame, bg='white')
        button_frame.pack(fill='x', padx=20, pady=20)
        
        # 重置配置按钮
        reset_btn = tk.Button(button_frame, text="重置为默认设置", 
                             font=('微软雅黑', 10),
                             bg='#e74c3c', fg='white', relief='flat',
                             command=self.reset_config)
        reset_btn.pack(side='left', padx=10)
        
        # 导出配置按钮
        export_btn = tk.Button(button_frame, text="导出配置", 
                              font=('微软雅黑', 10),
                              bg='#3498db', fg='white', relief='flat',
                              command=self.export_config)
        export_btn.pack(side='left', padx=10)
        
        # 导入配置按钮
        import_btn = tk.Button(button_frame, text="导入配置", 
                              font=('微软雅黑', 10),
                              bg='#27ae60', fg='white', relief='flat',
                              command=self.import_config)
        import_btn.pack(side='left', padx=10)
        
        canvas.pack(side="left", fill="both", expand=True)
        scrollbar.pack(side="right", fill="y")
        
    def setup_about_tab(self):
        """关于选项卡"""
        about_frame = ttk.Frame(self.notebook)
        self.notebook.add(about_frame, text="关于")
        
        about_text = """
        🖥️ 桌面助手 v1.0
        
        功能特色：
        • 一键整理桌面文件
        • 智能分类管理
        • 安全备份还原
        • 本地化存储
        
        开发理念：
        参考360桌面助手、元气桌面等知名软件
        专注于桌面整理的核心功能
        
        使用说明：
        1. 点击"扫描桌面"查看当前桌面文件
        2. 点击"一键整理"自动分类整理文件
        3. 点击"还原桌面"可恢复原始状态
        4. 点击"打开收纳空间"查看整理结果
        
        注意事项：
        • 系统文件和重要快捷方式不会被移动
        • 所有操作都有备份记录，可安全还原
        • 建议定期清理收纳空间
        """
        
        about_label = tk.Label(about_frame, text=about_text, font=('微软雅黑', 11), 
                              justify='left', bg='white')
        about_label.pack(fill='both', expand=True, padx=20, pady=20)
        
    def scan_desktop(self):
        """扫描桌面文件"""
        try:
            desktop_files = []
            for item in os.listdir(self.desktop_path):
                item_path = os.path.join(self.desktop_path, item)
                if os.path.isfile(item_path):
                    desktop_files.append(item)
            
            # 按类型分类统计
            categorized = {}
            for category in self.file_categories:
                categorized[category] = []
            categorized['其他'] = []
            
            for file in desktop_files:
                file_ext = os.path.splitext(file)[1].lower()
                categorized_flag = False
                for category, extensions in self.file_categories.items():
                    if file_ext in extensions:
                        categorized[category].append(file)
                        categorized_flag = True
                        break
                if not categorized_flag:
                    categorized['其他'].append(file)
            
            # 更新状态显示
            self.status_text.config(state='normal')
            self.status_text.delete(1.0, tk.END)
            
            total_files = len(desktop_files)
            status_info = f"桌面扫描完成 - 共发现 {total_files} 个文件\n"
            status_info += f"桌面路径: {self.desktop_path}\n\n"
            status_info += "文件分类统计:\n"
            status_info += "-" * 40 + "\n"
            
            for category, files in categorized.items():
                if files:
                    status_info += f"{category}: {len(files)} 个文件\n"
                    for file in files[:3]:  # 只显示前3个文件
                        status_info += f"  • {file}\n"
                    if len(files) > 3:
                        status_info += f"  • ... 还有 {len(files) - 3} 个文件\n"
                    status_info += "\n"
            
            self.status_text.insert(1.0, status_info)
            self.status_text.config(state='disabled')
            
            # 更新收纳空间概览
            self.update_storage_overview()
            
        except Exception as e:
            messagebox.showerror("错误", f"扫描桌面时出错: {str(e)}")
    
    def organize_desktop(self):
        """一键整理桌面"""
        confirm_needed = get_setting("confirm_before_organize", True)
        
        if not confirm_needed or messagebox.askyesno("确认整理", "确定要整理桌面吗？\n\n整理后的文件将移动到收纳空间，\n您可以随时还原。"):
            threading.Thread(target=self._organize_desktop_thread, daemon=True).start()
    
    def _organize_desktop_thread(self):
        """在后台线程中执行整理操作"""
        try:
            # 获取桌面文件
            desktop_files = []
            for item in os.listdir(self.desktop_path):
                item_path = os.path.join(self.desktop_path, item)
                if os.path.isfile(item_path):
                    desktop_files.append(item)
            
            if not desktop_files:
                messagebox.showinfo("提示", "桌面没有需要整理的文件！")
                return
            
            # 创建备份记录
            backup_record = {
                'timestamp': datetime.now().isoformat(),
                'files': [],
                'source': self.desktop_path
            }
            
            total_files = len(desktop_files)
            processed = 0
            
            # 按类型整理文件
            for file in desktop_files:
                file_path = os.path.join(self.desktop_path, file)
                file_ext = os.path.splitext(file)[1].lower()
                
                # 跳过系统文件和重要快捷方式（如果启用了此设置）
                if get_setting("skip_system_files", True) and self._is_system_file(file):
                    continue
                
                # 确定文件分类
                target_category = '其他'
                for category, extensions in self.file_categories.items():
                    if file_ext in extensions:
                        target_category = category
                        break
                
                # 创建目标目录
                target_dir = os.path.join(self.storage_path, target_category)
                os.makedirs(target_dir, exist_ok=True)
                
                # 移动文件
                target_path = os.path.join(target_dir, file)
                # 如果目标文件已存在，添加时间戳
                if os.path.exists(target_path):
                    name, ext = os.path.splitext(file)
                    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
                    new_name = f"{name}_{timestamp}{ext}"
                    target_path = os.path.join(target_dir, new_name)
                
                shutil.move(file_path, target_path)
                
                # 记录备份信息
                backup_record['files'].append({
                    'original_path': file_path,
                    'new_path': target_path,
                    'category': target_category
                })
                
                processed += 1
                self.progress_var.set((processed / total_files) * 100)
                self.root.update_idletasks()
            
            # 保存备份记录
            backup_file = os.path.join(self.backup_path, f"backup_{datetime.now().strftime('%Y%m%d_%H%M%S')}.json")
            with open(backup_file, 'w', encoding='utf-8') as f:
                json.dump(backup_record, f, ensure_ascii=False, indent=2)
            
            self.progress_var.set(0)
            messagebox.showinfo("整理完成", f"成功整理了 {processed} 个文件！\n\n您可以点击'打开收纳空间'查看结果。")
            
            # 刷新显示
            self.scan_desktop()
            
        except Exception as e:
            messagebox.showerror("错误", f"整理桌面时出错: {str(e)}")
            self.progress_var.set(0)
    
    def _is_system_file(self, filename):
        """判断是否为系统文件或重要快捷方式"""
        system_files = ['desktop.ini', 'thumbs.db', 'Thumbs.db']
        important_shortcuts = ['回收站', '此电脑', '网络', '控制面板']
        
        if filename in system_files:
            return True
        
        # 检查是否为重要系统快捷方式
        for important in important_shortcuts:
            if important in filename:
                return True
        
        return False
    
    def restore_desktop(self):
        """还原桌面"""
        try:
            # 获取备份文件列表
            backup_files = []
            for file in os.listdir(self.backup_path):
                if file.startswith('backup_') and file.endswith('.json'):
                    backup_files.append(file)
            
            if not backup_files:
                messagebox.showinfo("提示", "没有找到备份记录！")
                return
            
            # 选择最新的备份文件
            backup_files.sort(reverse=True)
            latest_backup = os.path.join(self.backup_path, backup_files[0])
            
            if messagebox.askyesno("确认还原", f"确定要还原到最近一次整理前的状态吗？\n\n备份时间: {backup_files[0].replace('backup_', '').replace('.json', '')}"):
                with open(latest_backup, 'r', encoding='utf-8') as f:
                    backup_data = json.load(f)
                
                restored_count = 0
                for file_info in backup_data['files']:
                    if os.path.exists(file_info['new_path']):
                        shutil.move(file_info['new_path'], file_info['original_path'])
                        restored_count += 1
                
                messagebox.showinfo("还原完成", f"成功还原了 {restored_count} 个文件！")
                self.scan_desktop()
                
        except Exception as e:
            messagebox.showerror("错误", f"还原桌面时出错: {str(e)}")
    
    def open_storage(self):
        """打开收纳空间"""
        try:
            os.startfile(self.storage_path)
        except Exception as e:
            messagebox.showerror("错误", f"无法打开收纳空间: {str(e)}")
    
    def update_storage_overview(self):
        """更新收纳空间概览"""
        try:
            # 清空现有数据
            for item in self.storage_tree.get_children():
                self.storage_tree.delete(item)
            
            if not os.path.exists(self.storage_path):
                return
            
            # 统计每个分类的文件
            for category in self.file_categories.keys():
                category_path = os.path.join(self.storage_path, category)
                if os.path.exists(category_path):
                    files = os.listdir(category_path)
                    file_count = len(files)
                    
                    # 计算总大小
                    total_size = 0
                    for file in files:
                        file_path = os.path.join(category_path, file)
                        if os.path.isfile(file_path):
                            total_size += os.path.getsize(file_path)
                    
                    size_str = self._format_size(total_size)
                    
                    if file_count > 0:
                        self.storage_tree.insert('', 'end', text=f"📁 {category}", 
                                               values=(file_count, size_str))
            
            # 检查其他分类
            other_path = os.path.join(self.storage_path, '其他')
            if os.path.exists(other_path):
                files = os.listdir(other_path)
                if files:
                    file_count = len(files)
                    total_size = sum(os.path.getsize(os.path.join(other_path, f)) 
                                   for f in files if os.path.isfile(os.path.join(other_path, f)))
                    size_str = self._format_size(total_size)
                    self.storage_tree.insert('', 'end', text="📁 其他", 
                                           values=(file_count, size_str))
                    
        except Exception as e:
            print(f"更新收纳空间概览时出错: {str(e)}")
    
    def _format_size(self, size_bytes):
        """格式化文件大小"""
        if size_bytes == 0:
            return "0 B"
        
        units = ['B', 'KB', 'MB', 'GB']
        unit_index = 0
        size = float(size_bytes)
        
        while size >= 1024 and unit_index < len(units) - 1:
            size /= 1024
            unit_index += 1
        
        return f"{size:.1f} {units[unit_index]}"
    
    def on_setting_changed(self):
        """设置更改时的回调"""
        try:
            # 保存设置
            set_setting("auto_backup", self.auto_backup_var.get())
            set_setting("confirm_before_organize", self.confirm_organize_var.get())
            set_setting("skip_system_files", self.skip_system_var.get())
            set_setting("max_backup_count", int(self.backup_count_var.get()))
        except Exception as e:
            print(f"保存设置时出错: {e}")
    
    def reset_config(self):
        """重置配置为默认值"""
        if messagebox.askyesno("确认重置", "确定要重置所有设置为默认值吗？\n这将清除您的自定义配置。"):
            try:
                from config import config
                config.reset_to_default()
                messagebox.showinfo("重置完成", "配置已重置为默认值，请重启程序以应用更改。")
            except Exception as e:
                messagebox.showerror("错误", f"重置配置时出错: {str(e)}")
    
    def export_config(self):
        """导出配置"""
        try:
            filepath = filedialog.asksaveasfilename(
                title="导出配置文件",
                defaultextension=".json",
                filetypes=[("JSON文件", "*.json"), ("所有文件", "*.*")]
            )
            if filepath:
                from config import config
                if config.export_config(filepath):
                    messagebox.showinfo("导出成功", f"配置已导出到:\n{filepath}")
                else:
                    messagebox.showerror("导出失败", "导出配置文件时出现错误")
        except Exception as e:
            messagebox.showerror("错误", f"导出配置时出错: {str(e)}")
    
    def import_config(self):
        """导入配置"""
        try:
            filepath = filedialog.askopenfilename(
                title="导入配置文件",
                filetypes=[("JSON文件", "*.json"), ("所有文件", "*.*")]
            )
            if filepath:
                from config import config
                if config.import_config(filepath):
                    messagebox.showinfo("导入成功", "配置已导入，请重启程序以应用更改。")
                else:
                    messagebox.showerror("导入失败", "导入配置文件时出现错误\n请检查文件格式是否正确")
        except Exception as e:
            messagebox.showerror("错误", f"导入配置时出错: {str(e)}")
    
    def run(self):
        """运行应用程序"""
        self.root.mainloop()

if __name__ == "__main__":
    app = DesktopHelper()
    app.run() 